import asyncio
import os
import time

import dotenv
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()

os.environ['OPENAI_API_KEY'] = os.getenv('OPENAI_API_KEY')
os.environ['OPENAI_BASE_URL'] = os.getenv('OPENAI_BASE_URL')

CHAT_MODEL = ChatOpenAI(model="gpt-4o-mini")

# 同步调用
def sync_test():
    message1 = [SystemMessage(content="你是一位乐于助人的智能小助手"), HumanMessage(content="请帮我介绍什么是机器学习")]
    start_time = time.time()
    response = CHAT_MODEL.invoke(message1)
    duration = time.time() - start_time
    print(f"同步调用耗时:{duration:.2f}秒")
    return response, duration

# 异步调用
async def async_test():
    message1 = [SystemMessage(content="你是一位乐于助手的智能小助手"), HumanMessage(content="请帮我介绍什么是机器学习")]
    start_time = time.time()
    response = await CHAT_MODEL.ainvoke(message1)   # 异步调用
    duration = time.time() - start_time
    print(f"异步调用耗时:{duration:.2f}秒")
    return response, duration

if __name__ == "__main__":
    # 运行同步测试
    sync_response, sync_duration = sync_test()
    print(f"同步响应内容:{sync_response.content[:100]}...\n")
    # 运行异步测试
    async_response, async_duration = asyncio.run(async_test())
    print(f"异步响应内容:{async_response.content[:100]}...\n")

    print("\n=== 并发测试 ===")
    start_time = time.time()

    async def run_concurrent_tests():
        # 创建3个异步任务
        tasks = [async_test() for _ in range(3)]
        # 并发执行所有任务
        return await asyncio.gather(*tasks)

    # 执行并发测试
    results = asyncio.run(run_concurrent_tests())

    total_time = time.time() - start_time
    print(f"\n3个并发异步调用总耗时:{total_time:.2f}秒")
    print(f"\n平均每个调用耗时：{total_time / 3:.2f}秒")